import os
import cv2
import numpy as np


class ImageProcessor:
    def __init__(self):
        pass

    def imwrite_unicode(self, path, image):
        """保存图像，支持Unicode路径"""
        ext = '.' + path.split('.')[-1]
        success, encoded_image = cv2.imencode(ext, image)
        if success:
            encoded_image.tofile(path)
            return True
        return False

    def imread_unicode(self, path, color=cv2.IMREAD_COLOR):
        data = np.fromfile(path, dtype=np.uint8)
        img = cv2.imdecode(data, color)
        return img

    def crop_black_border_ignore_text(self, image_path):
        save_path = image_path
        img = self.imread_unicode(image_path)
        h, w, _ = img.shape

        # 转灰度用于边界识别
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # 创建非黑掩码：排除背景太暗的像素
        _, binary = cv2.threshold(gray, 20, 255, cv2.THRESH_BINARY)

        # 进一步排除文字（边缘小区域）：只保留大块区域
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 合并大轮廓区域（例如主体内容）
        mask = np.zeros_like(gray)
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > 5000:  # 忽略小块区域（可能是文字）
                cv2.drawContours(mask, [cnt], -1, 255, -1)

        # 再次查找非黑区域边界（这次不包含文字）
        coords = cv2.findNonZero(mask)
        x, y, bw, bh = cv2.boundingRect(coords)

        # 进一步收缩上下边界（如果还存在边框残留）
        margin = 5
        top_crop = max(y - margin, 0)
        bottom_crop = min(y + bh + margin, h)

        cropped_img = img[top_crop:bottom_crop, x:x + bw]

        self.imwrite_unicode(save_path, cropped_img)
        return cropped_img

    def process_edge_images_in_folder(self,input_dir):
        for filename in os.listdir(input_dir):
            if not filename.lower().endswith(('.jpg', '.png', '.jpeg', '.bmp')):
                continue

            filepath = os.path.join(input_dir, filename)
            image = self.imread_unicode(filepath)
            if image is None:
                continue

            basename = os.path.splitext(filename)[0]

            # Edge detection (Laplacian)
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            blurred = cv2.GaussianBlur(gray, (5, 5), 0)
            laplacian = cv2.Laplacian(blurred, cv2.CV_64F)
            sharpness = np.uint8(np.absolute(laplacian))
            _, edge_mask = cv2.threshold(sharpness, 20, 255, cv2.THRESH_BINARY)
            self.imwrite_unicode(os.path.join(input_dir, f"{basename}_edge.png"), edge_mask)

    def enhance_for_ocr(self, img, aggressive=False):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        enhanced = cv2.equalizeHist(gray)
        _, binary = cv2.threshold(enhanced, 180, 255, cv2.THRESH_BINARY)

        if aggressive:
            edge = cv2.Laplacian(binary, cv2.CV_8U)
            result = cv2.addWeighted(binary, 0.8, edge, 0.2, 0)
            return result
        return binary
